import tensorflow as tf 
import numpy as np
import pathlib
from tensorflow import keras
from matplotlib import pyplot as plt
def load_from_path(all_image_paths, num, save_path_train, save_path_valid):
    '''读取所有图片'''
    image_count = len(all_image_paths)
    for i in range(0, image_count):
        image = tf.io.read_file(all_image_paths[i])
        image = tf.image.decode_jpeg(image)
        image = tf.image.adjust_gamma(image, 0.6)  # Gamma 
        image = tf.image.resize_with_crop_or_pad(image, 100,100)
        image_augmentation(image, num, save_path_train+str(i),save_path_valid+str(i))
        print(i,"------->crop end")

def image_augmentation(image, num , save_path_train, save_path_valid):
    '''图片裁剪增强'''
    # (1470-64+1) * (1780-64+1) = 1407 * 1717
    crop_y = [i for i in range(37)]
    crop_x = [i for i in range(37)]
    index_x = np.random.choice(crop_x, num-1, replace=False)
    index_y = np.random.choice(crop_y, num-1, replace=False)
    image1 = tf.image.resize_with_crop_or_pad(image, 64, 64)
    img = tf.image.encode_jpeg(image1) # jpeg
    # 保存图片
    with tf.io.gfile.GFile(save_path_train + "-0"+'.jpeg','wb') as file:
        file.write(img.numpy())

    for j in range(num-1):
        crop_image = tf.image.crop_to_bounding_box(image, index_y[j], index_x[j], 64, 64)
        if j < 2:
            path = save_path_train+"-"+str(j+1)+'.jpeg'
        else:
            path = save_path_valid+"-"+str(j-2)+'.jpeg'
        '''
        if j < 9 :
            path = save_path_train+"-00"+str(j+1)+'.jpeg'
        elif j < 39:
            path = save_path_train+"-0"+str(j+1)+'.jpeg'
        elif j < 49:
            path = save_path_valid+"-00"+str(j-39)+'.jpeg'
        else:
            path = save_path_valid+"-0"+str(j-39)+'.jpeg'
        '''
        crop_image = tf.image.encode_jpeg(crop_image) # jpeg
        with tf.io.gfile.GFile(path,'wb') as file:
            file.write(crop_image.numpy())    



def load_clutter_from_path(all_image_paths, num, save_path_train, save_path_valid):
    '''读取所有图片'''
    image_count = len(all_image_paths)
    for i in range(0, image_count):
        image = tf.io.read_file(all_image_paths[i])
        image = tf.image.decode_jpeg(image)
        image = tf.image.adjust_gamma(image, 0.6)  # Gamma 
        image = tf.image.resize_with_crop_or_pad(image, 1780,1470)
        image_augmentation(image, num, save_path_train+str(i),save_path_valid+str(i))
        print(i,"------->crop end")

def image_clutter_augmentation(image, num , save_path_train, save_path_valid):
    '''图片裁剪增强'''
    # (1470-64+1) * (1780-64+1) = 1407 * 1717
    crop_y = [i for i in range(1717)]
    crop_x = [i for i in range(1407)]
    index_x = np.random.choice(crop_x, num-1, replace=False)
    index_y = np.random.choice(crop_y, num-1, replace=False)
    image1 = tf.image.resize_with_crop_or_pad(image, 64, 64)
    img = tf.image.encode_jpeg(image1) # jpeg
    # 保存图片
    with tf.io.gfile.GFile(save_path_train + "-000"+'.jpeg','wb') as file:
        file.write(img.numpy())

    for j in range(num-1):
        crop_image = tf.image.crop_to_bounding_box(image, index_y[j], index_x[j], 64, 64)
        
        if j < 9 :
            path = save_path_train+"-00"+str(j+1)+'.jpeg'
        elif j < 39:
            path = save_path_train+"-0"+str(j+1)+'.jpeg'
        elif j < 49:
            path = save_path_valid+"-00"+str(j-39)+'.jpeg'
        else:
            path = save_path_valid+"-0"+str(j-39)+'.jpeg'
        crop_image = tf.image.encode_jpeg(crop_image) # jpeg
        with tf.io.gfile.GFile(path,'wb') as file:
            file.write(crop_image.numpy())  

def get_paths(path, label):
    '''获取数据集路径'''
    # 获得数据集文件路径
    data_path = pathlib.Path(path)
    # 获得所有类别图片的路径
    all_image_paths = list(data_path.glob('*/*'))
    all_image_paths = [str(path1) for path1 in all_image_paths]
    # 数据集图片数量
    image_count = len(all_image_paths)
    print(image_count)
    for image in all_image_paths[:5] : 
        print(image, ' --->  ', label)
    all_image_labels = [label for i in range(image_count)]
    return all_image_paths, all_image_labels
    


def main():
    '''主函数'''
    all_image_paths, all_image_labels = get_paths('E:/大学课程资料/大四上/毕业设计/MSTAR_Clutter/Target', 1)
    load_from_path(all_image_paths, 4, 'E:/大学课程资料/大四上/毕业设计/MSTAR_Clutter/Train/Target/', 'E:/大学课程资料/大四上/毕业设计/MSTAR_Clutter/Valid/Target/')

    
    #all_image_paths, all_image_labels = get_paths('E:/大学课程资料/大四上\毕业设计/MSTAR官方数据/data/MSTAR_PUBLIC_CLUTTER_CD2/Clutter2', 0)
    #load_from_path(all_image_paths, 60, 'E:/大学课程资料/大四上/毕业设计/MSTAR_Clutter/Train/Clutter/', 'E:/大学课程资料/大四上/毕业设计/MSTAR_Clutter/Valid/Clutter/')


if __name__ == '__main__':
    main()